import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

import sys
sys.path.append("../Hartree-Fock/")
from hartreefock import *


def VMC_energy(mol: Molecule) -> float:
    """
    使用蒙特卡洛计算某一个分子Hartree-Fock波函数的能量
    参数：
        mol: 分子
    返回值：
        能量
    """
    EHF, __, basis, Cx = HartreeFock(mol, details=True)
    C = Cx[:, 0]

    Nn = len(mol.atoms) # 原子数目
    Nb = len(basis)  # 基矢维度
    Ne = -mol.charge # 电子数目
    for atom in mol.atoms:
        Ne += atom.Z
    assert Ne % 2 == 0

    # 原子核间势能
    En = 0
    for i in range(Nn):
        for j in range(i+1, Nn):
            En += mol.atoms[i].Z * mol.atoms[j].Z / \
                abs(mol.atoms[i].R - mol.atoms[j].R)

    # 初始化电子坐标
    R = []
    for i in range(Ne):
        R.append(np.random.normal(scale = 0.5, size = 3))
    
    E = []
    Ek_arr = []
    Eext_arr = []
    Ee_arr = []

    # total_step = 50 * 10000 # 50w
    total_step = 10000
    for step in range(total_step):
        if step % 1000 == 0:
            print(f"{step} step")
        # 动能项
        Ek = 0
        for ri in R:
            de = 0
            nu = 0
            for i in range(Nb):
                s = basis[i]
                z = ri[2]
                r = np.linalg.norm(ri[:2])
                de += C[i] * eval_sto3g(s, z, r)
                for j in range(3):
                    g = s.g[j]
                    R2 = r**2 + (z - g.center)**2
                    nu += C[i] * s.d[j] * (3*g.alpha - 2 * g.alpha**2 * R2) * eval_gaussian(g, z, r)
            Ek += nu / de

        # 原子核吸引势
        Eext = 0
        for atom in mol.atoms:
            for ri in R:
                RiI = np.linalg.norm([ri[0], ri[1], ri[2] - atom.R])
                Eext += -atom.Z / RiI
        
        # 电子排斥势
        Ee = 0
        for i in range(Ne):
            for j in range(i+1, Ne):
                Ee += 1 / np.linalg.norm(R[i] - R[j])
        
        E.append(Ek + Eext + Ee + En)
        Ek_arr.append(Ek)
        Eext_arr.append(Eext)
        Ee_arr.append(Ee)

        delta = np.random.normal(scale = 0.1, size = 3)
        change_index = np.random.randint(Ne)
        r1 = R[change_index]
        r2 = R[change_index] + delta
        psi1 = 0
        psi2 = 0
        for i in range(Nb):
            s = basis[i]
            z = r1[2]
            r = np.linalg.norm(r1[:2])
            psi1 += C[i] * eval_sto3g(s, z, r)

            z = r2[2]
            r = np.linalg.norm(r2[:2])
            psi2 += C[i] * eval_sto3g(s, z, r)
        
        P = min(1, (psi2/psi1)**2)
        if np.random.rand() < P:
            R[change_index] = r2
    

    meanE = np.cumsum(E) / np.array(range(1, len(E)+1))
    plt.xlabel("step")
    plt.ylabel("Energy/a.u.")
    plt.plot(range(len(E)), E, linewidth=0.2, label="VMC")
    plt.plot([0, len(E)], [EHF, EHF], '--', linewidth='0.5', label="RHF")
    plt.legend()
    plt.savefig("instant-energy.png")

    plt.clf()
    plt.xlabel("step")
    plt.ylabel("Energy/a.u.")
    plt.plot(range(len(E)), meanE, linewidth=0.5, label="VMC")
    plt.plot([0, len(E)], [EHF, EHF], '--', linewidth='0.5', label="RHF")
    plt.legend()
    plt.savefig("mean-energy.png")



    # block analysis
    Nb = 20
    Lb = len(E) // Nb
    block = np.zeros((Nb,))
    for i in range(Nb):
        block[i] = np.mean(E[i*Lb:(i+1)*Lb])
    plt.clf()
    plt.xlim([0, Nb])
    plt.xlabel("block")
    plt.xticks(range(0, Nb+1, 2))
    plt.ylabel("Energy/a.u.")
    plt.plot(range(1, Nb+1), block, 'o')
    plt.savefig('block-anasis.png')

    v = np.var(block, ddof=1) / Nb
    s = np.sqrt(v)
    m = np.mean(E)
    print(f"VMC energy = {m:.5f} ± {s:.5f} a.u.")
    print(f"RHF energy = {EHF:.5f} a.u.")

if __name__ == '__main__':
    VMC_energy(H2(1.3))